# app/rag/vector_store_factory.py

from ..core.config import settings
from ..core.logger import logger
from pymilvus import connections, utility

# 尝试导入 Milvus 相关模块
try:
    from langchain_milvus import Milvus
    MILVUS_AVAILABLE = True
except ImportError:
    MILVUS_AVAILABLE = False
    logger.warning("Milvus 依赖未安装，将无法使用 Milvus 向量数据库")

# 尝试导入 Chroma 相关模块
try:
    from langchain_chroma import Chroma
    CHROMA_AVAILABLE = True
except ImportError:
    CHROMA_AVAILABLE = False
    logger.warning("Chroma 依赖未安装，将无法使用 Chroma 向量数据库")

from .embedding.base import EmbeddingProvider


class VectorStoreFactory:
    """
    向量数据库工厂类，根据配置创建相应的向量数据库实例
    """
    
    @staticmethod
    def create_vector_store(embedding_provider: EmbeddingProvider):
        """
        根据配置创建向量数据库实例
        
        Args:
            embedding_provider: 嵌入提供者实例
            
        Returns:
            向量数据库实例
        """
        vector_store_type = settings.VECTOR_STORE_TYPE
        
        if vector_store_type == "milvus" and MILVUS_AVAILABLE:
            logger.info("初始化 Milvus 向量数据库")
            return VectorStoreFactory._create_milvus_store(embedding_provider)
            
        elif vector_store_type == "chroma" and CHROMA_AVAILABLE:
            logger.info("初始化 Chroma 向量数据库")
            return VectorStoreFactory._create_chroma_store(embedding_provider)
            
        else:
            # 默认使用 Chroma
            if CHROMA_AVAILABLE:
                logger.info("默认使用 Chroma 向量数据库")
                return VectorStoreFactory._create_chroma_store(embedding_provider)
            elif MILVUS_AVAILABLE:
                logger.info("默认使用 Milvus 向量数据库")
                return VectorStoreFactory._create_milvus_store(embedding_provider)
            else:
                raise RuntimeError("没有可用的向量数据库，请安装 Milvus 或 Chroma 依赖")

    @staticmethod
    def _create_milvus_store(embedding_provider: EmbeddingProvider):
        """创建 Milvus 向量数据库实例"""
        try:
            # 检查 Milvus 服务连接
            connections.connect(
                alias="default", 
                host=settings.MILVUS_HOST, 
                port=settings.MILVUS_PORT
            )
            
            # 检查集合是否存在
            collection_exists = utility.has_collection(settings.MILVUS_COLLECTION_NAME)
            logger.info(f"Milvus 集合 '{settings.MILVUS_COLLECTION_NAME}' 是否存在: {collection_exists}")
            
            # 创建或连接到 Milvus 集合
            # 注意：对于langchain-milvus==0.1.6版本，我们不能直接指定vector_dim
            # 我们需要确保嵌入函数输出的向量维度与集合期望的维度匹配
            vector_store = Milvus(
                embedding_function=embedding_provider.get_embedding(),
                connection_args={
                    "host": settings.MILVUS_HOST,
                    "port": settings.MILVUS_PORT
                },
                collection_name=settings.MILVUS_COLLECTION_NAME,
                # 如果集合不存在，会自动创建
                drop_old=False,  # 不删除已存在的集合
                auto_id=True,  # 自动生成ID，避免警告
            )
            
            logger.info("成功初始化 Milvus 向量数据库")
            return vector_store
            
        except Exception as e:
            logger.error(f"初始化 Milvus 向量数据库失败: {e}")
            raise e

    @staticmethod
    def _create_chroma_store(embedding_provider: EmbeddingProvider):
        """创建 Chroma 向量数据库实例"""
        try:
            # 确保目录存在
            import os
            if not os.path.exists(settings.CHROMA_PERSIST_DIRECTORY):
                os.makedirs(settings.CHROMA_PERSIST_DIRECTORY, exist_ok=True)
            
            vector_store = Chroma(
                collection_name=settings.CHROMA_COLLECTION_NAME,
                embedding_function=embedding_provider.get_embedding(),
                persist_directory=settings.CHROMA_PERSIST_DIRECTORY,
            )
            
            logger.info("成功初始化 Chroma 向量数据库")
            return vector_store
            
        except Exception as e:
            logger.error(f"初始化 Chroma 向量数据库失败: {e}")
            raise e